package common;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 消息对象 - 用于封装和解析协议消息
 */
public class Message {
    private byte version;       // 协议版本
    private byte type;          // 消息类型
    private int requestId;      // 请求ID
    private byte[] data;        // 消息数据
    private Map<String, String> headers = new ConcurrentHashMap<>();  // 消息头

    // 请求ID生成器
    private static final AtomicInteger idGenerator = new AtomicInteger(1);

    /**
     * 创建一个新的消息
     */
    public Message(byte type, byte[] data) {
        this.version = Protocol.PROTOCOL_VERSION;
        this.type = type;
        this.requestId = idGenerator.getAndIncrement();
        this.data = data != null ? data : new byte[0];
    }

    /**
     * 创建一个带请求ID的新消息
     */
    public Message(byte type, int requestId, byte[] data) {
        this.version = Protocol.PROTOCOL_VERSION;
        this.type = type;
        this.requestId = requestId;
        this.data = data != null ? data : new byte[0];
    }

    /**
     * 解码ByteBuffer为消息对象
     */
    public static Message decode(ByteBuffer buffer) {
        byte version = buffer.get();
        byte type = buffer.get();
        int requestId = buffer.getInt();
        int length = buffer.getInt();

        byte[] data = new byte[length];
        buffer.get(data);

        return new Message(type, requestId, data);
    }

    /**
     * 将消息编码为ByteBuffer
     */
    public ByteBuffer encode() {
        ByteBuffer buffer = ByteBuffer.allocate(Protocol.HEADER_SIZE + data.length);
        buffer.put(version);
        buffer.put(type);
        buffer.putInt(requestId);
        buffer.putInt(data.length);
        buffer.put(data);
        buffer.flip();
        return buffer;
    }

    // Getters and setters
    public byte getVersion() {
        return version;
    }

    public byte getType() {
        return type;
    }

    public int getRequestId() {
        return requestId;
    }

    public byte[] getData() {
        return data;
    }

    public Map<String, String> getHeaders() {
        return headers;
    }

    public void addHeader(String key, String value) {
        headers.put(key, value);
    }

    public String getHeader(String key) {
        return headers.get(key);
    }

    /**
     * 创建心跳消息
     */
    public static Message heartbeat() {
        return new Message(Protocol.MSG_TYPE_HEARTBEAT, null);
    }

    /**
     * 创建认证消息
     */
    public static Message auth(String token) {
        return new Message(Protocol.MSG_TYPE_AUTH, token.getBytes(StandardCharsets.UTF_8));
    }

    /**
     * 创建新连接消息
     */
    public static Message newConnection(int connectionId) {
        ByteBuffer buffer = ByteBuffer.allocate(4);
        buffer.putInt(connectionId);
        return new Message(Protocol.MSG_TYPE_NEW_CONN, buffer.array());
    }

    /**
     * 创建数据消息
     */
    public static Message data(int connectionId, byte[] data) {
        ByteBuffer buffer = ByteBuffer.allocate(4 + data.length);
        buffer.putInt(connectionId);
        buffer.put(data);
        buffer.flip();
        byte[] messageData = new byte[buffer.remaining()];
        buffer.get(messageData);
        return new Message(Protocol.MSG_TYPE_DATA, messageData);
    }

    /**
     * 创建关闭连接消息
     */
    public static Message closeConnection(int connectionId) {
        ByteBuffer buffer = ByteBuffer.allocate(4);
        buffer.putInt(connectionId);
        return new Message(Protocol.MSG_TYPE_CLOSE_CONN, buffer.array());
    }
}
